# -*- coding: utf-8 -*-

import os
import duckdb


def convert_csv_query(query, column_delimiter=','):
    front_statement = query[:query.upper().find("FROM")]
    if query.upper().find("WHERE") != -1:
        table_name = query[query.upper().find("FROM") + 4:query.upper().find("WHERE"):]
        end_statement = query[query.upper().find("WHERE"):]
    elif query.upper().find("GROUP BY") != -1:
        table_name = query[query.upper().find("FROM") + 4:query.upper().find("GROUP BY"):]
        end_statement = query[query.upper().find("GROUP BY"):]
    elif query.upper().find("ORDER BY") != -1:
        table_name = query[query.upper().find("FROM") + 4:query.upper().find("ORDER BY"):]
        end_statement = query[query.upper().find("ORDER BY"):]
    elif query.upper().find("LIMIT") != -1:
        table_name = query[query.upper().find("FROM") + 4:query.upper().find("LIMIT"):]
        end_statement = query[query.upper().find("LIMIT"):]
    else:
        table_name = query[query.upper().find("FROM") + 4:]
        end_statement = ''
    new_query = f"{front_statement} FROM read_csv_auto({table_name}, delim='{column_delimiter}') {end_statement}"
    count_query = f"SELECT COUNT(*) FROM {table_name}"
    return new_query, count_query


def query_data_by_duckdb(query, cursor=None):
    try:
        duckdb_value = cursor.sql(query)
        header = duckdb_value.columns
        data = duckdb_value.fetchall()
        return header, data
    finally:
        cursor.close()


def load_file_by_duckdb(cursor, file_path, table_name, is_csv_file=True, column_delimiter=',',
                        row_delimiter='\r\n'):
    try:
        duckdb_data = duckdb.read_csv(file_path, delimiter=column_delimiter) if is_csv_file else duckdb.read_parquet(
            file_path)

        columns = duckdb_data.columns
        dtypes = [dtype.id for dtype in duckdb_data.dtypes]  # duckdb_data.dtypes
        header_and_datatypes = list(zip(columns, dtypes))
        # 创建表（如果不存在）
        columns = ', '.join(
            [f"\"{header[0]}\" {get_table_column_type(header[1])}" for header in header_and_datatypes])
        create_table_query = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns})"
        cursor.execute(create_table_query)
        if is_csv_file:
            copy_table_query = f"COPY {table_name} FROM '{file_path}' (FORMAT CSV, HEADER, DELIMITER '{column_delimiter}')"
        else:
            copy_table_query = f"COPY {table_name} FROM '{file_path}' (FORMAT PARQUET)"
        cursor.execute(copy_table_query)
    finally:
        pass


def get_table_column_type(pd_type: str):
    # if pd_type.find('int') >= 0:
    #     return 'INTEGER'
    # elif pd_type.find('float') >= 0:
    #     return 'REAL'
    # elif pd_type.find('decimal') >= 0:
    #     return 'NUMERIC'
    # else:
    return 'TEXT'


def filename_to_table_name(filename):
    # 去除文件扩展名
    # base_name = os.path.splitext(filename)[0]
    base_name = filename
    # 转换为小写
    base_name = base_name.lower()

    # 替换特殊字符
    table_name = ''.join(c if c.isalnum() else '_' for c in base_name)

    # 可选：添加前缀或后缀
    # table_name = f"fu_{table_name}"

    return table_name